Mixture test

# Packages
  library(tidyverse)
  library(tidybayes)
  library(rethinking)
  library(patchwork)
  library(cmdstanr)
  library(GGally)
  library(posterior)
  library(future)
  library(parallel)

# Functions
  post_plot <- readRDS("./functions/post_plot.rds")
  post_plot_2 <- readRDS("./functions/post_plot_2.rds")
  plot_interval <- readRDS("./functions/plot_interval.rds")
  hist_plot <- readRDS("./functions/hist_plot.rds")
  binom_pmf <- readRDS("./functions/binom_pmf.rds")
  hist_plot_simple <- readRDS("./functions/hist_plot_simple.rds")

1 Generative model

1.1 Fixed parameters

The intercepts \(\alpha\).

 a <- tibble(
   "1" = c(log(3.8 * 12 * 60), log(3.6 * 12 * 60), log(3.5 * 12 * 60), log(3.2 * 12 * 60)),
   "2" = c(log(4.2), log(4.5), log(4.8), log(5)),
   "gr" = c(log(4.8), log(5), log(5.2), log(5.8)),
   "1_gr" = c(-0.65, -0.7, -0.6, 0.65),
   "2_gr" = c(-0.8, -1, -0.9, -0.85),
   "gr_1" = c(1.9, 2, 2.1, 1.85),
   "gr_2" = c(1.4, 1.6, 1.8, 1.9)
 )

The effects of sex, \(\mathbf{\Omega}^{1}\), \(\mathbf{\Omega}^{2}\), and \(\mathbf{\Omega}^{\text{gr}}\).

# State 1
Om_1 <- matrix(c(0.275, 0.075, 0.075, -0.425), nrow = 2, byrow = FALSE)

# State 2
Om_2 <- matrix(c(-0.4, 0, 0, 0.4), nrow = 2, byrow = FALSE)

# State 3 or 4
Om_gr <- matrix(c(-0.325, 0.175, -0.225, 0.425), nrow = 2, byrow = FALSE)

# List of Omegas
sex_effects <- list(Om_1, Om_2, Om_gr)

The effect of \(U_{[a]}^{\text{ind}}\) on \(\phi\).

b_phi <- tibble(
  "1" = -0.5,
  "2" = 0.5,
  "give_gr" = 0.5,
  "rec_gr" = 0.5,
  "1_give_gr" = 0.1,
  "1_rec_gr" = 0.1, 
  "2_give_gr" = 0.5,
  "2_rec_gr" = 0.5,
  "give_gr_1" = -0.5,
  "give_gr_2" = -0.5,
  "rec_gr_1" = -0.5,
  "rec_gr_2" = -0.5
)

The effect of \(U_{[a, b]}^{\text{dyad}}\) (where \(U_{[a, b]}^{\text{dyad}} = U_{[b, a]}^{\text{dyad}}\)) on \(\tau\).

b_tau <- tibble(
  "1" = -0.5,
  "2" = 0.5,
  "gr" = 0.5,
  "1_gr" = 0.5,
  "2_gr" = 0.5,
  "gr_1" = -0.5,
  "gr_2" = -0.5
)

1.2 stochastic_param

stochastic_param <- function(
  N_ind_vec = c(3, 4, 3),
  N_group = length(N_ind_vec),
  N_dyad_vec = ((N_ind_vec * N_ind_vec) - N_ind_vec) / 2,
  alpha = a,
  Omega = sex_effects,
  beta_phi = b_phi,
  beta_tau = b_tau,
  sigma_U_ind = 0.5,
  sigma_U_dyad = 0.5,
  sigma_ups_ind = 0.1,
  sigma_ups_dyad = 0.1) {
  
  params <- list()
  
  for (group in 1:N_group) {
    # Nb of ind and dyads in the group
    N_ind <- N_ind_vec[group]
    N_dyad <- N_dyad_vec[group]
    
    ## 1. Individual features
    ID_features <- tibble(
      ID = c(1:N_ind),
      z = sample(1:2, N_ind, replace = TRUE)
    )
    
    ## 2. all dyads
    dyads <- tibble(
      ind_a = t(combn(N_ind, 2))[, 1],
      ind_b = t(combn(N_ind, 2))[, 2],
      dyad_id = c(1:N_dyad)
    ) %>%
      # Add a's attribute
      left_join(ID_features, by = c("ind_a" = "ID")) %>%
      rename(z_a = z) %>%
      
      # Add b's attribute
      left_join(ID_features, by = c("ind_b" = "ID")) %>%
      rename(z_b = z)
    
    ## 2. Exogenous variation
    # U
    U_ind <- rnorm(N_ind, 0, sigma_U_ind)
    U_dyad <- rnorm(N_dyad, 0, sigma_U_dyad)
    
    # upsilon individual
    ups_ind <- tibble(
      "1" = rnorm(N_ind, 0, sigma_ups_ind),
      "2" = rnorm(N_ind, 0, sigma_ups_ind),
      "give_gr" = rnorm(N_ind, 0, sigma_ups_ind),
      "rec_gr" = rnorm(N_ind, 0, sigma_ups_ind),
      "1_give_gr" = rnorm(N_ind, 0, sigma_ups_ind),
      "1_rec_gr" = rnorm(N_ind, 0, sigma_ups_ind),
      "2_give_gr" = rnorm(N_ind, 0, sigma_ups_ind),
      "2_rec_gr" = rnorm(N_ind, 0, sigma_ups_ind),
      "give_gr_1" = rnorm(N_ind, 0, sigma_ups_ind),
      "give_gr_2" = rnorm(N_ind, 0, sigma_ups_ind),
      "rec_gr_1" = rnorm(N_ind, 0, sigma_ups_ind),
      "rec_gr_2" = rnorm(N_ind, 0, sigma_ups_ind),
    )
    
    # upsilon dyad
    ups_dyad <- tibble(
      "1" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "2" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "gr_ab" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "gr_ba" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "1_gr_ab" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "1_gr_ba" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "2_gr_ab" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "2_gr_ba" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "gr_1_ab" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "gr_2_ab" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "gr_1_ba" = rnorm(N_dyad, 0, sigma_ups_dyad),
      "gr_2_ba" = rnorm(N_dyad, 0, sigma_ups_dyad),
    )
    
    ## 3. phi
    phi <- tibble(
      "1" = beta_phi$`1` * U_ind + ups_ind$`1`,
      "2" = beta_phi$`2` * U_ind + ups_ind$`2`,
      "give_gr" = beta_phi$`give_gr` * U_ind + ups_ind$`give_gr`,
      "rec_gr" = beta_phi$`rec_gr` * U_ind + ups_ind$`rec_gr`,
      "1_give_gr" = beta_phi$`1_give_gr` * U_ind + ups_ind$`1_give_gr`,
      "1_rec_gr" = beta_phi$`1_rec_gr` * U_ind + ups_ind$`1_rec_gr`,
      "2_give_gr" = beta_phi$`2_give_gr` * U_ind + ups_ind$`2_give_gr`,
      "2_rec_gr" = beta_phi$`2_rec_gr` * U_ind + ups_ind$`2_rec_gr`,
      "give_gr_1" = beta_phi$`give_gr_1` * U_ind + ups_ind$`give_gr_1`,
      "give_gr_2" = beta_phi$`give_gr_2` * U_ind + ups_ind$`give_gr_2`,
      "rec_gr_1" = beta_phi$`rec_gr_1` * U_ind + ups_ind$`rec_gr_1`,
      "rec_gr_2" = beta_phi$`rec_gr_2` * U_ind + ups_ind$`rec_gr_2`,
    )
    
    ## 4. tau
    tau <- tibble(
      "1" = beta_tau$`1` * U_dyad + ups_dyad$`1`,
      "2" = beta_tau$`2` * U_dyad + ups_dyad$`2`,
      "gr_ab" = beta_tau$`gr` * U_dyad + ups_dyad$`gr_ab`,
      "gr_ba" = beta_tau$`gr` * U_dyad + ups_dyad$`gr_ba`,
      "1_gr_ab" = beta_tau$`1_gr` * U_dyad + ups_dyad$`1_gr_ab`,
      "1_gr_ba" = beta_tau$`1_gr` * U_dyad + ups_dyad$`1_gr_ba`,
      "2_gr_ab" = beta_tau$`2_gr` * U_dyad + ups_dyad$`2_gr_ab`,
      "2_gr_ba" = beta_tau$`2_gr` * U_dyad + ups_dyad$`2_gr_ba`,
      "gr_1_ab" = beta_tau$`gr_1` * U_dyad + ups_dyad$`gr_1_ab`,
      "gr_2_ab" = beta_tau$`gr_2` * U_dyad + ups_dyad$`gr_2_ab`,
      "gr_1_ba" = beta_tau$`gr_1` * U_dyad + ups_dyad$`gr_1_ba`,
      "gr_2_ba" = beta_tau$`gr_2` * U_dyad + ups_dyad$`gr_2_ba`,
    )
    
    
    ## 5. theta
    theta <- array(data = NA, dim = c(N_dyad, 4))
    for (dyad in 1:N_dyad) {
      theta[dyad, 1] <- exp(alpha$`1`[group] +
                            Omega[[1]][dyads$z_a[dyad], dyads$z_b[dyad]] +
                            phi$`1`[dyads$ind_a[dyad]] +
                            phi$`1`[dyads$ind_b[dyad]] +
                            tau$`1`[dyad])
      
      theta[dyad, 2] <- exp(alpha$`2`[group] +
                            Omega[[2]][dyads$z_a[dyad], dyads$z_b[dyad]] +
                            phi$`2`[dyads$ind_a[dyad]] +
                            phi$`2`[dyads$ind_b[dyad]] +
                            tau$`2`[dyad])
      
      theta[dyad, 3] <- exp(alpha$`gr`[group] +
                            Omega[[3]][dyads$z_a[dyad], dyads$z_b[dyad]] +
                            phi$`give_gr`[dyads$ind_a[dyad]] +
                            phi$`rec_gr`[dyads$ind_b[dyad]] +
                            tau$gr_ab[dyad])
      
      theta[dyad, 4] <- exp(alpha$`gr`[group] +
                            Omega[[3]][dyads$z_b[dyad], dyads$z_a[dyad]] +
                            phi$`rec_gr`[dyads$ind_a[dyad]] +
                            phi$`give_gr`[dyads$ind_b[dyad]] +
                            tau$`gr_ba`[dyad])
    } # dyad
    
    ## 5.Psi
    Psi <- array(data = NA, dim = c(4, 4, N_dyad))
    for (dyad in 1:N_dyad) {
      # Firs row
     Psi[1, 2, dyad] <- 0
     Psi[1, 3, dyad] <- alpha$`1_gr`[group] +
        phi$`1_give_gr`[dyads$ind_a[dyad]] +
        phi$`1_rec_gr`[dyads$ind_b[dyad]] +
        tau$`1_gr_ab`[dyad]
     Psi[1, 4, dyad] <- alpha$`1_gr`[group] +
        phi$`1_rec_gr`[dyads$ind_a[dyad]] +
        phi$`1_give_gr`[dyads$ind_b[dyad]] +
        tau$`1_gr_ba`[dyad]
      
      # Second row
     Psi[2, 1, dyad] <- 0
     Psi[2, 3, dyad] <- alpha$`2_gr`[group] +
        phi$`2_give_gr`[dyads$ind_a[dyad]] +
        phi$`2_rec_gr`[dyads$ind_b[dyad]] +
        tau$`2_gr_ab`[dyad]
     Psi[2, 4, dyad] <- alpha$`2_gr`[group] +
        phi$`2_rec_gr`[dyads$ind_a[dyad]] +
        phi$`2_give_gr`[dyads$ind_b[dyad]] +
        tau$`2_gr_ba`[dyad]
      
      # Third row
      Psi[3, 1, dyad] <- alpha$`gr_1`[group] +
        phi$`give_gr_1`[dyads$ind_a[dyad]] +
        phi$`rec_gr_1`[dyads$ind_b[dyad]] +
        tau$`gr_1_ab`[dyad]
      Psi[3, 2, dyad] <- alpha$`gr_2`[group] +
        phi$`give_gr_2`[dyads$ind_a[dyad]] +
        phi$`rec_gr_2`[dyads$ind_b[dyad]] +
        tau$`gr_2_ab`[dyad]
      Psi[3, 4, dyad] <- 0
      
      # Fourth row
      Psi[4, 1, dyad] <- alpha$`gr_1`[group] +
        phi$`give_gr_1`[dyads$ind_b[dyad]] +
        phi$`rec_gr_1`[dyads$ind_a[dyad]] +
        tau$`gr_1_ba`[dyad]
      Psi[4, 2, dyad] <- alpha$`gr_2`[group] +
        phi$`give_gr_2`[dyads$ind_b[dyad]] +
        phi$`rec_gr_2`[dyads$ind_a[dyad]] +
        tau$`gr_2_ba`[dyad]
      Psi[4, 3, dyad] <- 0
    } # ind
    
    ## 6. Gamma
    Gamma <- array(data = NA, dim = c(4, 4, N_dyad))
    for (dyad in 1:N_dyad) {
      for (k in 1:4) {
        # Zero diagonal
        Gamma[k, k, dyad] <- 0
      } # end for k
      
      # Row 1
      for (k in 2:4) {
        Gamma[1, k, dyad] <- exp(Psi[1, k, dyad]) /
          sum(exp(Psi[1, , dyad]), na.rm = T)
      } # end for k
      
      # Row 2
      for (k in c(1, 3, 4)) {
        Gamma[2, k, dyad] <- exp(Psi[2, k, dyad]) /
          sum(exp(Psi[2, , dyad]), na.rm = T)
      } # end for k
      
      # Row 3
      for (k in c(1, 2, 4)) {
        Gamma[3, k, dyad] <- exp(Psi[3, k, dyad]) /
          sum(exp(Psi[3, , dyad]), na.rm = T)
      } # end for k
      
      # Row 4
      for (k in 1:3) {
        Gamma[4, k, dyad] <- exp(Psi[4, k, dyad]) /
          sum(exp(Psi[4, , dyad]), na.rm = T)
      } # end for k
    } # end for dyad
    
    params[[group]] <- list(
      dyads = dyads,
      U_ind = U_ind,
      U_dyad = U_dyad,
      ups_ind = ups_ind,
      ups_dyad = ups_dyad,
      phi = phi,
      tau = tau,
      theta = theta,
      Psi = Psi,
      Gamma = Gamma
    )
  } # end for group
  
  params %>% return()
} # end stochastic_param

1.3 ctmc

ctmc <- function(N_ind_vec = c(3, 3, 4),
                 N_groups = length(N_ind_vec),
                 N_dyad_vec = ((N_ind_vec * N_ind_vec) - N_ind_vec) / 2,
                 params,
                 delta = 5,
                 end_sim = c(1e3, 1e3, 1e3),
                 ...) {
  
  # First level of df is group
  df <- list()
  
  for (group in 1:N_groups) {
    # Second level of df is dyad
    df[[group]] <- list()
    N_dyad <- N_dyad_vec[group]
    
    for (dyad in 1:N_dyad) {
      ## 1. First sojourn (state 1)
      (df[[group]][[dyad]] <- tibble(
        id = c(1, 1),          # observation id
        long = rbinom(1, 1, 0.5) %>% rep(2),
        x = rexp(1, 1 / (params[[group]]$theta[dyad, 1])) %>% rep(2), # long holding time
        t = c(0, x[1]),        # time stamps
        s = c(1, 1),           # state
        z = c("start", "end"), # modifier
        c = 0                  # censoring level
      ))
      
      ## 2. Run the CTMC
      repeat {
        # Next transition
        last_state <- df[[group]][[dyad]]$s %>% last()
        new_state <- rmultinom(1, 1, params[[group]]$Gamma[,,dyad][last_state,]) %>%
          which.max()
        
        # Next sojourn
        df[[group]][[dyad]] <- df[[group]][[dyad]] %>%
          add_row(
            id = df[[group]][[dyad]]$id %>% last() + 1,
            long = rbinom(1, 1, 0.5),
            x = ifelse(new_state == 1 & long == 0,
                       rexp(1, 1 / delta), # short holding time in s = 1
                       rexp(1, 1 / (params[[group]]$theta[dyad, new_state])) 
                       ),
            t = df[[group]][[dyad]]$t %>% last(),
            s = new_state,
            z = "start",
            c = 0
          )
        df[[group]][[dyad]] <- df[[group]][[dyad]] %>%
          add_row(
            id = df[[group]][[dyad]]$id %>% last(),
            long = df[[group]][[dyad]]$long %>% last(),
            x = df[[group]][[dyad]]$x %>% last(),
            t = (df[[group]][[dyad]]$t %>% last()) +
              last(df[[group]][[dyad]]$x),
            s = df[[group]][[dyad]]$s %>% last(),
            z = "end",
            c = 0
          )
        
        # break if we're after the maximal time
        if (tail(df[[group]][[dyad]]$t, 1) >= end_sim[group]) {
          break
        } # end if
      } # end break
    } # end for dyad
  } # end for group
  df %>% return()
} # end ctmc

For example:

set.seed(2666)
params <- stochastic_param(N_ind_vec = c(3, 4))
(seq_states_ex <- ctmc(N_ind_vec = c(3, 4), params = params))

[[1]]
[[1]][[1]]
# A tibble: 4 × 7
     id  long        x     t     s z         c
  <dbl> <int>    <dbl> <dbl> <dbl> <chr> <dbl>
1     1     0 1222.       0      1 start     0
2     1     0 1222.    1222.     1 end       0
3     2     0    0.295 1222.     3 start     0
4     2     0    0.295 1223.     3 end       0

[[1]][[2]]
# A tibble: 4 × 7
     id  long        x     t     s z         c
  <dbl> <int>    <dbl> <dbl> <dbl> <chr> <dbl>
1     1     0 3577.       0      1 start     0
2     1     0 3577.    3577.     1 end       0
3     2     1    0.881 3577.     2 start     0
4     2     1    0.881 3578.     2 end       0

[[1]][[3]]
# A tibble: 4 × 7
     id  long       x     t     s z         c
  <dbl> <int>   <dbl> <dbl> <dbl> <chr> <dbl>
1     1     0 2680.      0      1 start     0
2     1     0 2680.   2680.     1 end       0
3     2     0    1.23 2680.     4 start     0
4     2     0    1.23 2681.     4 end       0


[[2]]
[[2]][[1]]
# A tibble: 4 × 7
     id  long       x     t     s z         c
  <dbl> <int>   <dbl> <dbl> <dbl> <chr> <dbl>
1     1     0 2774.      0      1 start     0
2     1     0 2774.   2774.     1 end       0
3     2     1    1.94 2774.     3 start     0
4     2     1    1.94 2776.     3 end       0

[[2]][[2]]
# A tibble: 4 × 7
     id  long        x     t     s z         c
  <dbl> <int>    <dbl> <dbl> <dbl> <chr> <dbl>
1     1     1 1222.       0      1 start     0
2     1     1 1222.    1222.     1 end       0
3     2     1    0.919 1222.     2 start     0
4     2     1    0.919 1223.     2 end       0

[[2]][[3]]
# A tibble: 4 × 7
     id  long        x      t     s z         c
  <dbl> <int>    <dbl>  <dbl> <dbl> <chr> <dbl>
1     1     0 17199.       0      1 start     0
2     1     0 17199.   17199.     1 end       0
3     2     1     5.05 17199.     2 start     0
4     2     1     5.05 17204.     2 end       0

[[2]][[4]]
# A tibble: 10 × 7
      id  long       x      t     s z         c
   <dbl> <int>   <dbl>  <dbl> <dbl> <chr> <dbl>
 1     1     1   92.8     0       1 start     0
 2     1     1   92.8    92.8     1 end       0
 3     2     1    8.19   92.8     4 start     0
 4     2     1    8.19  101.      4 end       0
 5     3     0    2.52  101.      2 start     0
 6     3     0    2.52  103.      2 end       0
 7     4     1    3.80  103.      3 start     0
 8     4     1    3.80  107.      3 end       0
 9     5     1 3643.    107.      1 start     0
10     5     1 3643.   3750.      1 end       0

[[2]][[5]]
# A tibble: 16 × 7
      id  long       x     t     s z         c
   <dbl> <int>   <dbl> <dbl> <dbl> <chr> <dbl>
 1     1     1  165.      0      1 start     0
 2     1     1  165.    165.     1 end       0
 3     2     0   15.8   165.     2 start     0
 4     2     0   15.8   181.     2 end       0
 5     3     0   11.2   181.     1 start     0
 6     3     0   11.2   192.     1 end       0
 7     4     1    8.87  192.     2 start     0
 8     4     1    8.87  201.     2 end       0
 9     5     1    4.65  201.     4 start     0
10     5     1    4.65  206.     4 end       0
11     6     0   22.5   206.     1 start     0
12     6     0   22.5   228.     1 end       0
13     7     1    7.76  228.     4 start     0
14     7     1    7.76  236.     4 end       0
15     8     1 1852.    236.     1 start     0
16     8     1 1852.   2088.     1 end       0

[[2]][[6]]
# A tibble: 10 × 7
      id  long        x     t     s z         c
   <dbl> <int>    <dbl> <dbl> <dbl> <chr> <dbl>
 1     1     0  240.       0      1 start     0
 2     1     0  240.     240.     1 end       0
 3     2     1    1.25   240.     3 start     0
 4     2     1    1.25   241.     3 end       0
 5     3     1    0.302  241.     4 start     0
 6     3     1    0.302  242.     4 end       0
 7     4     1    0.751  242.     2 start     0
 8     4     1    0.751  242.     2 end       0
 9     5     1 2179.     242.     1 start     0
10     5     1 2179.    2422.     1 end       0

1.4 d_ff

d_ff <- function(df, start_focal, end_focal){
for (dyad in 1:length(df)){
# Censor the start of the dyadic focal follow 
# i.e. we cut anything before start_focal 
  df[[dyad]] <- df[[dyad]] %>%
  add_row(t = start_focal, c = 1) %>%
  add_row(t = end_focal, c = 2) %>%
  arrange(t)  %>%
  mutate(
    id = case_when(
      c == 1 ~ lag(id),
      c == 2 ~ lead(id),
      TRUE ~ id
    ),
    long = case_when(
      c == 1 ~ lag(long),
      c == 2 ~ lead(long),
      TRUE ~ long
    ),
    s = case_when(
      c == 1 ~ lag(s),
      c == 2 ~ lead(s),
      TRUE ~ s
    ),
    z = case_when(
      c == 1 ~ "start",
      c == 2 ~ "end",
      TRUE ~ z
    ),
    x = case_when(
      c == 1 ~ lead(t) - start_focal,
      c == 2 ~ end_focal - lag(t),
      TRUE ~ x
    )
  ) %>% group_by(id) %>%
  mutate(
    x = case_when(
      sum(c) == 0 ~ x,                                       
      sum(c) %in% c(1, 2) ~ x[2],
      sum(c) == 3 ~ end_focal - start_focal                           
    )
  ) %>%
  ungroup() %>% 
  filter(t >= start_focal & t <= end_focal)
  } # end for i
  
  return(df)
} # end d_ff

Suppose I focal individual 1 of group 2, thereby essentially focaling the dyads 1, 2, and 3 of group 2:

seq_states_ex[[2]][1:3] %>% d_ff(start_focal = 320, end_focal = 360)

[[1]]
# A tibble: 2 × 7
     id  long     x     t     s z         c
  <dbl> <int> <dbl> <dbl> <dbl> <chr> <dbl>
1     1     0    40   320     1 start     1
2     1     0    40   360     1 end       2

[[2]]
# A tibble: 2 × 7
     id  long     x     t     s z         c
  <dbl> <int> <dbl> <dbl> <dbl> <chr> <dbl>
1     1     1    40   320     1 start     1
2     1     1    40   360     1 end       2

[[3]]
# A tibble: 2 × 7
     id  long     x     t     s z         c
  <dbl> <int> <dbl> <dbl> <dbl> <chr> <dbl>
1     1     0    40   320     1 start     1
2     1     0    40   360     1 end       2

1.5 sim

sim <- function(N_ind_vec,
                N_group = length(N_ind_vec),
                N_dyad_vec = ((N_ind_vec * N_ind_vec) - N_ind_vec) / 2,
                parameters = params,
                n_focals_vec,
                burnin = 1e3,
                duration_focal = 40,
                end_sim_vec = burnin + n_focals_vec * duration_focal) {
  ## 1. Define a number of objects
  dyads <- list()     # list of dyads and individuals per group
  sim_data <- list()  # simulated observations (sojourns in states per dyad)
  d <- list()         # summary of sim_data
  
  ## 2. We run ctmc for each group
  state_seq <- ctmc(N_ind_vec = N_ind_vec,
                    params = parameters,
                    end_sim = end_sim_vec)
  
  for (group in 1:N_group) {
    ## 3. Define indidividuals & dyads
    # We give individuals unique name across groups
    # e.g. 1:10 (group 1), 11:26 (group 2), etc.
    bl_ind <- c(0, head(N_ind_vec, -1))[1:group] %>% sum()
    bl_dyad <- c(0, head(N_dyad_vec, -1))[1:group] %>% sum()
    
    N_ind <- N_ind_vec[group]
    dyads[[group]] <- tibble(
      grp = group,
      ind_a = bl_ind + t(combn(N_ind, 2))[, 1],
      ind_b = bl_ind + t(combn(N_ind, 2))[, 2],
      z_a = parameters[[group]]$dyads$z_a,
      z_b = parameters[[group]]$dyads$z_b,
      dyad_id = bl_dyad + c(1:N_dyad_vec[group])
    )
    
    
    ## 4. Define sequence of focal follows for each group
    # 4.1 Sequence of individuals ids to focal by group
    # We make sure to focal each individual at least once
    focal_ind <- bl_ind + c(1:N_ind)
    
    # For the remaining focal-follows, we draw individuals at random
    focal_ind[(N_ind + 1):n_focals_vec[group]] <-
      sample(x = bl_ind + 1:N_ind,
             size = n_focals_vec[group] - N_ind,
             replace = TRUE)
    
    # 4.2 We mark the start time of the focal-follows
    time_stamps <- seq(from = burnin, to = end_sim_vec[[group]], by = duration_focal)
    
    dyads_to_sample <- list()    # dyads to sample for each focal follow
    sim_data[[group]] <- list()  # observed sequence of states
    d[[group]] <- list()
    
    # for each focal follow (e.g. ff of ind 1, of ind 3, of ind 4, ...)
    for (focal in 1:n_focals_vec[group]) {
      # sim_data is indexed by (i) group; (ii) ind. focal follow; (iii) dyad.
      sim_data[[group]][[focal]] <- list()
      d[[group]][[focal]] <- list()
      
      ## 5. We conduct the focal follows per se
      # Which *dyads* are we practically following for each ind. focal-follow
      dyads_to_sample[[focal]] <- dyads[[group]] %>%
        filter(ind_a == focal_ind[focal] |
                 ind_b == focal_ind[focal]) %>%
        pull(dyad_id)
      
      # We select the latent state sequence of interest
      sim_data[[group]][[focal]] <- 
        state_seq[[group]][dyads_to_sample[[focal]] - bl_dyad]
      
      # We cut the windows corresponding to focal-follows
      sim_data[[group]][[focal]] <- d_ff(sim_data[[group]][[focal]],
                                         time_stamps[focal],
                                         time_stamps[focal] + duration_focal)
      
   
      # for each dyad
      for (dyad in 1:(N_ind - 1)) {
        ## add name of focal individual and focal dyad for each
        sim_data[[group]][[focal]][[dyad]] <-
          sim_data[[group]][[focal]][[dyad]] %>%
          mutate(focal_id = focal_ind[focal],
                 focal_dyad = dyads_to_sample[[focal]][dyad])
        
        ## 6. Summarise each sojourn by a few variables
        d[[group]][[focal]][[dyad]] <-
          sim_data[[group]][[focal]][[dyad]] %>%
          group_by(id) %>%               # for each state state id
          summarize(
            grp = group,              # grp id
            long = first(long),       # long sub-state
            x = first(x),             # holding time
            s = first(s),             # state
            c = ifelse(sum(c) < 2, 0, 1),# censoring = {0, 1}
            dyad = first(focal_dyad), # id of dyad
            focal_id = first(focal_id), # id of focal individual
            prot_id = focal           # id of the protocol          
          ) %>%
          select(-id)
      } # end dyad
    } # end for focal
    
    d[[group]] <- d[[group]] %>% bind_rows()
  } # end for group
  
  ## 7. We collapse the sojourn data
  d <- d %>%
    bind_rows() %>%
    group_by(grp, focal_id, dyad, prot_id) %>%
    # add a unique identifier per dyadic focal follow
    mutate(prot_id = cur_group_id()) %>%
    ungroup() %>%
    arrange(prot_id)
  
  ## 8. Tansitions
  obs_tr <- tibble(
    grp = head(d$grp, -1),
    dyad_id = head(d$dyad, -1),
    prot_id_from = head(d$prot_id, -1),
    prot_id_to = tail(d$prot_id, -1),
    s_from = head(d$s, -1),
    s_to = tail(d$s, -1)
  ) %>%
    filter(prot_id_from == prot_id_to) %>%
    select(-c(prot_id_from, prot_id_to))
  
  ## 9. Exposure and "events" in each state
  exp_events <- d %>%
    mutate(observed = ifelse(c == 0, 1, 0)) %>%
    group_by(dyad, s) %>%
    # Amount of time spent in each state
    summarise(exposure = sum(x),
              # how many times have we observed the end?
              events = sum(observed)) %>%
    right_join(expand.grid(dyad = 1:sum(N_dyad_vec),
                           s = 1:4)) %>%
    arrange(dyad, s) %>%
    mutate_at(c("exposure", "events"), ~ replace_na(., 0))
    
  
  # Stan-friendly format
  list <- list()
  
  # General indices
  list$N_dyad <- sum(N_dyad_vec)
  list$N_ind <- sum(N_ind_vec)
  list$N_group <- N_group
  list$N_ind_per_group <- N_ind_vec
  list$N_dyad_per_group <- N_dyad_vec  
  
  # Individuals and dyads (dataset "A")
  list$A_grp <- bind_rows(dyads)$grp
  list$A_ind_a <- bind_rows(dyads)$ind_a
  list$A_ind_b <- bind_rows(dyads)$ind_b
  list$A_dyad <- bind_rows(dyads)$dyad_id
  list$A_z_a <- bind_rows(dyads)$z_a
  list$A_z_b <- bind_rows(dyads)$z_b
  
  # Observations j: sequence of states (dataset "B")
  list$J <- bind_rows(d) %>% nrow()
  list$B_grp <- d$grp
  list$B_x <- d$x
  list$B_s <- d$s
  list$B_c <- d$c
  list$B_long <- d$long
  list$B_dyad <- d$dyad
  
  # Observed transitions (dataset "C")
  list$N_trans <- nrow(obs_tr)
  list$C_grp <- obs_tr$grp
  list$C_dyad <- obs_tr$dyad_id
  list$C_s_from <- obs_tr$s_from
  list$C_s_to <- obs_tr$s_to
  
  # Observed exposure time and count of events in each state (dataset "D")
  list$N_row_D <- nrow(exp_events)
  list$D_dyad <- exp_events$dyad
  list$D_s <- exp_events$s
  list$D_exposure <- exp_events$exposure
  list$D_events <- exp_events$events

  list %>% return()  
} # end sim

1.6 Run sim

Function to call sim in parallel.

run_sim <- function(i, n_individuals, n_focals, param, 
                    duration_f = 40){  
  d <- sim(
    burnin = 1e3,
    duration_focal = duration_f,
    N_ind_vec = n_individuals,
    n_focals_vec = n_focals,
    parameters = param[[i]]
  )
  
  return(d)
}

We run it:

param_1 <- list() # stochastic parameters
d_1 <- list()     # simulated data

# We run sim 10 times
set.seed(546)
for (i in 1:10) {
  param_1[[i]] <- stochastic_param(N_ind_vec = c(10, 10, 10, 10))
}

param_1 %>% saveRDS("./sim_data/02/param_1.rds")

# we run the sim function (parallelised)
d_1 <- mclapply(
  1:10, 
  function(i) run_sim(i, 
                      n_individuals = c(10, 10, 10, 10),
                      n_focals = c(1500, 1500, 1500, 1500),
                      param = param_1)
)

# We save the data
  d_1 %>% saveRDS("./sim_data/02/d_1.rds")

2 Statistical model

# Stan model
  (m_1 <- cmdstan_model("./stan_models/extended_model.stan"))

data{
  // Indices
  int<lower = 1> N_dyad;  
  int<lower = 1> N_ind;  
  int<lower = 1> N_group;   
  
  // Dataset "A": all dyads, and corresponding grp and ind
  array[N_dyad] int A_dyad;        // Dyad id {1, ..., N_dyad}
  array[N_dyad] int A_grp;         // Group id {1, ..., N_group}
  array[N_dyad] int A_ind_a;       // ind a per dyad number
  array[N_dyad] int A_ind_b;       // ind b per dyad number
  array[N_dyad] int<lower = 1, upper = 2> A_z_a;  // sex of ind a
  array[N_dyad] int<lower = 1, upper = 2> A_z_b;  // sex of ind a

  // Dataset "B": observed states
  int<lower = 1> J;                // Number of observed states
  array[J] int<lower = 1, upper = N_group> B_grp; // Group
  vector<lower = 0>[J] B_x;                       // Holding time
  array[J] int<lower = 1, upper = 4> B_s;  // State {1, 2, 3, 4}
  array[J] int<lower = 0, upper = 1> B_c;  // censoring {0, 1}
  array[J] int<lower = 1, upper = N_dyad> B_dyad; // dyad {1, ..., N_dyad}

  // Dataset "C": observed transitions
  int<lower = 1> N_trans;          // Number of observed transitions
  array[N_trans] int<lower = 1, upper = N_dyad> C_dyad;  // dyad {1, ..., N_dyad}
  array[N_trans] int<lower = 1, upper = 4> C_s_from;     // state j
  array[N_trans] int<lower = 1, upper = 4> C_s_to;       // state j + 1
}

transformed data {
  // Count of transitions. One matrix per dyad
  array[N_dyad, 4, 4] int<lower = 0> transition_counts;

  // Initialise the transition matrices with zeros
  for (dyad in 1:N_dyad) {
    for (k in 1:4) {
      for (l in 1:4) {
        transition_counts[dyad, k, l] = 0;
      } // l
    } // k
  } // dyad

  // Fill the transition matrices with transitions from the input data
  for (tr in 1:N_trans) {
    transition_counts[C_dyad[tr], C_s_from[tr], C_s_to[tr]] += 1;
  } // tr
} // block

parameters {
  // Intercepts per group
  vector<lower = 4>[N_group] alpha_1;
  vector[N_group] alpha_2;
  vector[N_group] alpha_gr;
  vector[N_group] alpha_1_gr;
  vector[N_group] alpha_2_gr;
  vector[N_group] alpha_gr_1;
  vector[N_group] alpha_gr_2;
  
  // Offset for sex effects
  vector[10] Omega_raw;
  
  // Mixture parameters
  real<lower = 0, upper = 1> proba;
  real<lower = 0, upper = 60> delta;
  
  // Individual varying effects
  matrix[12, N_ind] z_phi; // z-score of phi (wide format: 12 R, N_ind C)
  vector<lower = 0>[12] sigma_phi; // SD of phi
  cholesky_factor_corr[12] L_phi; // Cholesky factor of corr. matrix phi
  
  // Dyadic varying effects
  matrix[12, N_dyad] z_tau; // z-score of tau (wide format: 12 R, N_dyad C)
  vector<lower = 0>[7] sigma_tau; // SD of tau
  cholesky_factor_corr[12] L_tau; // Cholesky factor of corr. matrix for tau
}

transformed parameters{
    // 1. One matrix per level for effect of sex
    // State 1
    matrix[2, 2] Omega_1;
    Omega_1[1, 1] = Omega_raw[1];
    Omega_1[2, 1] = Omega_raw[2];
    Omega_1[1, 2] = Omega_raw[2];
    Omega_1[2, 2] = Omega_raw[3];
    
    // State 2
    matrix[2, 2] Omega_2;
    Omega_2[1, 1] = Omega_raw[4];
    Omega_2[2, 1] = Omega_raw[5];
    Omega_2[1, 2] = Omega_raw[5];
    Omega_2[2, 2] = Omega_raw[6];
    
    // State 3
    matrix[2, 2] Omega_gr;
    Omega_gr[1, 1] = Omega_raw[7];
    Omega_gr[2, 1] = Omega_raw[8];
    Omega_gr[1, 2] = Omega_raw[9];
    Omega_gr[2, 2] = Omega_raw[10];
    
    // 2. individual varying effects
    matrix[N_ind, 12] phi_raw; // vertical: N_ind R, 12 C)
    phi_raw = (diag_pre_multiply(sigma_phi, L_phi) * z_phi)';
    vector[N_ind] phi_1 = phi_raw[, 1];
    vector[N_ind] phi_2 = phi_raw[, 2];
    vector[N_ind] phi_give_gr = phi_raw[, 3];
    vector[N_ind] phi_rec_gr = phi_raw[, 4];
    vector[N_ind] phi_1_give_gr = phi_raw[, 5];
    vector[N_ind] phi_1_rec_gr = phi_raw[, 6];
    vector[N_ind] phi_2_give_gr = phi_raw[, 7];
    vector[N_ind] phi_2_rec_gr = phi_raw[, 8];
    vector[N_ind] phi_give_gr_1 = phi_raw[, 9];
    vector[N_ind] phi_give_gr_2 = phi_raw[, 10];
    vector[N_ind] phi_rec_gr_1 = phi_raw[, 11];
    vector[N_ind] phi_rec_gr_2 = phi_raw[, 12];
    
  // 3. dyad varying effects
    matrix[N_dyad, 12] tau_raw; // vertical: N_ind R, 12 C
    tau_raw = (diag_pre_multiply([sigma_tau[1], sigma_tau[2], 
                                 sigma_tau[3], sigma_tau[3],
                                 sigma_tau[4], sigma_tau[4],
                                 sigma_tau[5], sigma_tau[5],
                                 sigma_tau[6], sigma_tau[7],
                                 sigma_tau[6], sigma_tau[7]], L_tau) * z_tau)';
    vector[N_dyad] tau_1 = tau_raw[, 1];
    vector[N_dyad] tau_2 = tau_raw[, 2];
    vector[N_dyad] tau_gr_ab = tau_raw[, 3];
    vector[N_dyad] tau_gr_ba = tau_raw[, 4];
    vector[N_dyad] tau_1_gr_ab = tau_raw[, 5];
    vector[N_dyad] tau_1_gr_ba = tau_raw[, 6];
    vector[N_dyad] tau_2_gr_ab = tau_raw[, 7];
    vector[N_dyad] tau_2_gr_ba = tau_raw[, 8];
    vector[N_dyad] tau_gr_1_ab = tau_raw[, 9];
    vector[N_dyad] tau_gr_2_ab = tau_raw[, 10];
    vector[N_dyad] tau_gr_1_ba = tau_raw[, 11];
    vector[N_dyad] tau_gr_2_ba = tau_raw[, 12];
  
  // We declare a bunch of local varibales that we define further below  
    array[N_dyad, 4] real theta;
    array[N_dyad, 4, 4] real Psi;
    array[N_dyad, 4, 4] real Gamma;  
    
  // 4. theta
    for (dyad in 1:N_dyad) {
    theta[dyad, 1] = exp(alpha_1[A_grp[dyad]] +
                     Omega_1[A_z_a[dyad], A_z_b[dyad]] +
                     phi_1[A_ind_a[dyad]] +
                     phi_1[A_ind_b[dyad]] +
                     tau_1[A_dyad[dyad]]);
    theta[dyad, 2] = exp(alpha_2[A_grp[dyad]] +
                     Omega_2[A_z_a[dyad], A_z_b[dyad]] +
                     phi_2[A_ind_a[dyad]] +
                     phi_2[A_ind_b[dyad]] +
                     tau_2[A_dyad[dyad]]);
    theta[dyad, 3] = exp(alpha_gr[A_grp[dyad]] +
                     Omega_gr[A_z_a[dyad], A_z_b[dyad]] +
                     phi_give_gr[A_ind_a[dyad]] +
                     phi_rec_gr[A_ind_b[dyad]] +
                     tau_gr_ab[A_dyad[dyad]]);
    theta[dyad, 4] = exp(alpha_gr[A_grp[dyad]] +
                     Omega_gr[A_z_b[dyad], A_z_a[dyad]] +
                     phi_give_gr[A_ind_b[dyad]] +
                     phi_rec_gr[A_ind_a[dyad]] +
                     tau_gr_ba[A_dyad[dyad]]);
     
  // 5. Psi                       
    Psi[dyad, 1, 2] = 0.0;
    Psi[dyad, 1, 3] = alpha_1_gr[A_grp[dyad]] + phi_1_give_gr[A_ind_a[dyad]] +
                             phi_1_rec_gr[A_ind_b[dyad]] + tau_1_gr_ab[A_dyad[dyad]];
    Psi[dyad, 1, 4] = alpha_1_gr[A_grp[dyad]] + phi_1_give_gr[A_ind_b[dyad]] +
                             phi_1_rec_gr[A_ind_a[dyad]] + tau_1_gr_ba[A_dyad[dyad]];
                             
    Psi[dyad, 2, 1] = 0.0;
    Psi[dyad, 2, 3] = alpha_2_gr[A_grp[dyad]] + phi_2_give_gr[A_ind_a[dyad]] +
                             phi_2_rec_gr[A_ind_b[dyad]] + tau_2_gr_ab[A_dyad[dyad]];
    Psi[dyad, 2, 4] = alpha_2_gr[A_grp[dyad]] + phi_2_give_gr[A_ind_b[dyad]] +
                             phi_2_rec_gr[A_ind_a[dyad]] + tau_2_gr_ba[A_dyad[dyad]];
                             
    Psi[dyad, 3, 1] = alpha_gr_1[A_grp[dyad]] + phi_give_gr_1[A_ind_a[dyad]] +
                             phi_rec_gr_1[A_ind_b[dyad]] + tau_gr_1_ab[A_dyad[dyad]];
    Psi[dyad, 3, 2] = alpha_gr_2[A_grp[dyad]] + phi_give_gr_2[A_ind_a[dyad]] +
                             phi_rec_gr_2[A_ind_b[dyad]] + tau_gr_2_ab[A_dyad[dyad]];
    Psi[dyad, 3, 4] = 0.0;
    
    Psi[dyad, 4, 1] = alpha_gr_1[A_grp[dyad]] + phi_give_gr_1[A_ind_b[dyad]] +
                             phi_rec_gr_1[A_ind_a[dyad]] + tau_gr_1_ba[A_dyad[dyad]];
    Psi[dyad, 4, 2] = alpha_gr_2[A_grp[dyad]] + phi_give_gr_2[A_ind_b[dyad]] +
                             phi_rec_gr_2[A_ind_a[dyad]] + tau_gr_2_ba[A_dyad[dyad]];
    Psi[dyad, 4, 3] = 0.0;

  // 6. Gamma
    for (k in 1:4) {
      real row_sum = 0.0;
  
      // Compute the denominator (sum of exponents in row k, excluding diagonal)
      for (m in 1:4) {
        if (m != k) {
          row_sum += exp(Psi[dyad, k, m]);
        } // end if
      } // end for m
  
      // Now compute Gamma for each element in row k
      for (l in 1:4) {
        if (k == l) {
          Gamma[dyad, k, l] = 0.0;  // Diagonal elements are 0
        } else {
          Gamma[dyad, k, l] = exp(Psi[dyad, k, l]) / row_sum;
        } // end if
      } // end for l
    } // end for k
  } // end for dyad
}

model{
  // 1. Likelihood for holding times
  for (j in 1:J){
        if (B_s[j] == 2 || B_s[j] == 3 || B_s[j] == 4){
          // not censored
          if (B_c[j] == 0){
            target += exponential_lpdf(B_x[j] | 1 / theta[B_dyad[j], B_s[j]]);
          } // end non-censored obs
        
          // censored
          if (B_c[j] == 1){
            target += exponential_lccdf(B_x[j] | 1 / theta[B_dyad[j], B_s[j]]);
          } // end censored obs
        } // end if state {2, 3, 4}
        
        if (B_s[j] == 1) {
          // not censored
          if (B_c[j] == 0){
            target += log_mix(proba,
                        exponential_lpdf(B_x[j] | 1.0 / theta[B_dyad[j], 1]),
                        exponential_lpdf(B_x[j] | 1.0 / delta)
                        );
          } // end non-censored obs
        
          // censored
          if (B_c[j] == 1){
            target += log_mix(proba,
                        exponential_lccdf(B_x[j] | 1.0 / theta[B_dyad[j], 1]),
                        exponential_lccdf(B_x[j] | 1.0 / delta)
                        );
          } // end censored obs
        } // end if state 1
    } // end for j
    
  // 2. Likelihood for observed transitions
  for (dyad in 1:N_dyad) {
    for (k in 1:4) { // Loop over current states k
      target += multinomial_lpmf(transition_counts[dyad, k] | 
                to_vector(Gamma[dyad, k]));
    } // k
  } // end for d
  
  // 3. Hyper-priors
    target += normal_lpdf(alpha_1 | 8, 2);
    target += normal_lpdf(alpha_2 | 1.5, 1);
    target += normal_lpdf(alpha_gr | 1.5, 1);
    target += normal_lpdf(delta | 10, 10);
    target += normal_lpdf(alpha_1_gr | 0, 1.5);
    target += normal_lpdf(alpha_2_gr | 0, 1.5);
    target += normal_lpdf(alpha_gr_1 | 0, 1.5);
    target += normal_lpdf(alpha_gr_2 | 0, 1.5);
    target += normal_lpdf(Omega_raw | 0, 1);
    target += beta_lpdf(proba | 5, 1);
    target += normal_lpdf(to_vector(z_phi) | 0, 1);
    target += normal_lpdf(to_vector(z_tau) | 0, 1);
    target += exponential_lpdf(sigma_phi | 1);
    target += exponential_lpdf(sigma_tau | 1);
    target += lkj_corr_cholesky_lpdf(L_phi | 4);
    target += lkj_corr_cholesky_lpdf(L_tau | 4);
} // end model block

generated quantities{
  // Corr. matrix from cholesky factors
    // Varying individual effects
    matrix[12, 12] c_ind;
    c_ind = multiply_lower_tri_self_transpose(L_phi);
  
    // Varying dyadic effects
    matrix[12, 12] c_dyad;
    c_dyad = multiply_lower_tri_self_transpose(L_tau);
                      
  // Marginal effect of sex
  vector[3] avg_s1; // average theta for 1_1, 1_2, 2_2 (S = 1)
  vector[3] avg_s2; // average theta for 1_1, 1_2, 2_2 (S = 2)
  vector[4] avg_gr; // average theta for 1_1, 1_2, 2_1, 2_2 (S = 3 or 4)
  
  vector[2] ATE_s1; // ATE for 1_2, 2_2 (S = 1)
  vector[2] ATE_s2; // ATE for 1_2, 2_2 (S = 2)
  vector[3] ATE_gr; // ATE for 1_2, 2_1 2_2 (S = 3 or 4)

  { // local scope
    array[N_dyad, 4] real theta_ns;     // thetas if there was no effect of sex
    array[N_dyad, 4] real theta_1_1;    // thetas if all dyads were male-male
    array[N_dyad, 4] real theta_1_2;    // thetas if all dyads were male-female
    array[N_dyad, 4] real theta_2_1;    // thetas if all dyads were female-male
    array[N_dyad, 4] real theta_2_2;    // thetas if all dyads were female-female
    for (dyad in 1:N_dyad){
    // state 1
    theta_ns[dyad, 1] = exp(alpha_1[A_grp[dyad]] +
                       phi_1[A_ind_a[dyad]] +
                       phi_1[A_ind_b[dyad]] +
                       tau_1[A_dyad[dyad]]);
    theta_1_1[dyad, 1] = theta_ns[dyad, 1] * exp(Omega_1[1, 1]);
    theta_1_2[dyad, 1] = theta_ns[dyad, 1] * exp(Omega_1[1, 2]);
    theta_2_2[dyad, 1] = theta_ns[dyad, 1] * exp(Omega_1[2, 2]);
    
    // state 2
    theta_ns[dyad, 2] = exp(alpha_2[A_grp[dyad]] +
                       phi_2[A_ind_a[dyad]] +
                       phi_2[A_ind_b[dyad]] +
                       tau_2[A_dyad[dyad]]);
    theta_1_1[dyad, 2] = theta_ns[dyad, 2] * exp(Omega_2[1, 1]);
    theta_1_2[dyad, 2] = theta_ns[dyad, 2] * exp(Omega_2[1, 2]);
    theta_2_2[dyad, 2] = theta_ns[dyad, 2] * exp(Omega_2[2, 2]);
    
    // state 3 and 4
    theta_ns[dyad, 3] = exp(alpha_gr[A_grp[dyad]] +
                       phi_give_gr[A_ind_a[dyad]] +
                       phi_rec_gr[A_ind_b[dyad]] +
                       tau_gr_ab[A_dyad[dyad]]);
    theta_1_1[dyad, 3] = theta_ns[dyad, 3] * exp(Omega_gr[1, 1]);
    theta_1_2[dyad, 3] = theta_ns[dyad, 3] * exp(Omega_gr[1, 2]);
    theta_2_1[dyad, 3] = theta_ns[dyad, 3] * exp(Omega_gr[2, 1]);
    theta_2_2[dyad, 3] = theta_ns[dyad, 3] * exp(Omega_gr[2, 2]);
    
    theta_ns[dyad, 4] = exp(alpha_gr[A_grp[dyad]] +
                       phi_give_gr[A_ind_b[dyad]] +
                       phi_rec_gr[A_ind_a[dyad]] +
                       tau_gr_ba[A_dyad[dyad]]);
    theta_1_1[dyad, 4] = theta_ns[dyad, 4] * exp(Omega_gr[1, 1]);
    theta_1_2[dyad, 4] = theta_ns[dyad, 4] * exp(Omega_gr[2, 1]);
    theta_2_1[dyad, 4] = theta_ns[dyad, 4] * exp(Omega_gr[1, 2]);
    theta_2_2[dyad, 4] = theta_ns[dyad, 4] * exp(Omega_gr[2, 2]);
    } // end for dyad
  
  // Average theta per state and sex combination
  avg_s1[1] = mean(theta_1_1[, 1]);
  avg_s1[2] = mean(theta_1_2[, 1]);
  avg_s1[3] = mean(theta_2_2[, 1]);
  avg_s2[1] = mean(theta_1_1[, 2]);
  avg_s2[2] = mean(theta_1_2[, 2]);
  avg_s2[3] = mean(theta_2_2[, 2]);
  avg_gr[1] = mean(append_row(
                      to_vector(theta_1_1[, 3]), to_vector(theta_1_1[, 4])));
  avg_gr[2] = mean(append_row(
                      to_vector(theta_1_2[, 3]), to_vector(theta_2_1[, 4])));
  avg_gr[3] = mean(append_row(
                      to_vector(theta_2_1[, 3]), to_vector(theta_1_2[, 4])));
  avg_gr[4] = mean(append_row(
                      to_vector(theta_2_2[, 3]), to_vector(theta_2_2[, 4])));
  
  // ATEs
  ATE_s1[1] = avg_s1[2] - avg_s1[1];
  ATE_s1[2] = avg_s1[3] - avg_s1[1];
  ATE_s2[1] = avg_s2[2] - avg_s2[1];
  ATE_s2[2] = avg_s2[3] - avg_s2[1];
  ATE_gr[1] = avg_gr[2] - avg_gr[1];
  ATE_gr[2] = avg_gr[3] - avg_gr[1];
  ATE_gr[3] = avg_gr[4] - avg_gr[1];
  } // end local variables
}

2.1 MCMC

Function to run the Stan model:

# Function to run Stan model and process posterior samples
run_stan_model <- function(i, 
                           d, 
                           m,
                           warmum = 1e3,
                           post_warmup = 1e3,
                           n_chains = 4,
                           para_chains = 4) {
  
  # Run Stan
  post <- m$sample(
    data = d[[i]],
    iter_warmup = warmum,
    iter_sampling = post_warmup,
    chains = n_chains,
    parallel_chains = para_chains,
    show_messages = TRUE
  )
  
  # Extract samples
  post_tidy <- post %>%
    tidy_draws() %>%
    mutate(iter = i)
  
  # Diagnostics
  diag <- post$summary() %>%
    mutate(iter = i)
  
  # diagnostics and post samples as output
  output <- list(
    diag = diag,
    samles = post_tidy
  )

  return(output)
}

We import the data:

# Import data
d_1 <- readRDS("./sim_data/02/d_1.rds")

We loop over the iterations:

options(warn = 1)
post_1 <- list()

# parallelise for loop
post_1 <- mclapply(
  1:10, 
  function(i) run_stan_model(i, d_1, m_1)
)

# Extract diagnostics from each iteration
diag_1 <- lapply(post_1, `[[`, "diag") %>% bind_rows()

# Extract posterior samples from all iterations and bind rows
post_1 <- lapply(post_1, `[[`, "samles") %>% bind_rows()

# We save the posterior
post_1 %>%
  saveRDS("./fitted_models/02.1/post_1.rds")
diag_1 %>%
  saveRDS("./fitted_models/02.1/diag_1.rds")

We import the posterior draws and diagnostics

post_1 <- readRDS("./fitted_models/02.1/post_1.rds")

2.2 Fixed effects

target <- a %>%
  mutate(grp = row_number()) %>%
  slice(1:4) %>%
  select(c(1:3, ncol(.))) %>%
  gather(param, value, 1:3) %>%
  mutate(value = value) %>%
  mutate(param = paste0("alpha_", param, "[", grp, "]")) %>%
  select(-grp) %>%
  mutate(iter = 1) %>%
  add_row(param = "delta", value = 5, iter = 1)

post_1 %>%
  select(starts_with("alpha_1["), starts_with("alpha_2["), starts_with("alpha_gr["), "delta", "iter") %>%
  gather(param, value, 1:(ncol(.)-1)) %>%
  post_plot(target,
            scale_slab = 0.5,
            alpha_filling = 0.1,
            min = 0, max = 15)

Warning: The `size` argument of `element_line()` is deprecated as of ggplot2 3.4.0.
ℹ Please use the `linewidth` argument instead.

Ignoring unknown labels:
• colour : "Percentile Interval"

Sex effects:

target_omega <- tibble(
  `Omega_1[1,1]` = sex_effects[[1]][1,1],
  `Omega_1[1,2]` = sex_effects[[1]][1,2],
  `Omega_1[2,2]` = sex_effects[[1]][2,2],
  `Omega_2[1,1]` = sex_effects[[2]][1,1],
  `Omega_2[1,2]` = sex_effects[[2]][1,2],
  `Omega_2[2,2]` = sex_effects[[2]][2,2],
  `Omega_gr[1,1]` = sex_effects[[3]][1,1],
  `Omega_gr[1,2]` = sex_effects[[3]][1,2],
  `Omega_gr[2,1]` = sex_effects[[3]][2,1],
  `Omega_gr[2,2]` = sex_effects[[3]][2,2]
) %>%
  gather(param, value, 1:(ncol(.))) %>%
  mutate(iter = 1)

post_1 %>%
  select(starts_with("Omega_1["), starts_with("Omega_2["), starts_with("Omega_gr["), "iter") %>%
  select(- c("Omega_1[2,1]", "Omega_2[2,1]")) %>%
  gather(param, value, 1:(ncol(.)-1)) %>%
  post_plot(target_omega,
            scale_slab = 0.5,
            min = -3, max = 3)

2.3 Posterior predictions

2.3.1 Holding times

d_1 <- readRDS("./sim_data/02/d_1.rds")

# Observed holding times for j in {1, ..., J}
  obs_ht <- tibble(
    dyad = d_1[[1]]$B_dyad,
    s = d_1[[1]]$B_s,
    x = d_1[[1]]$B_x,
    c = d_1[[1]]$B_c
  )

# theta posterior samples
thetas <- post_1 %>%
  filter(iter == 1) %>%
  select(starts_with("theta["), delta, proba) %>% 
  slice(1:30) %>%
  mutate(sample = 1:nrow(.)) %>%
  gather(param, value, 1:(ncol(.) - 3)) %>%
  extract(param, into = c("dyad", "s"), regex = "theta\\[(\\d+),(\\d+)\\]", convert = TRUE)

# For each observed holding time
post_pred_x <- list()
for (spl in 1:30) {
  post_pred_x[[spl]] <- obs_ht %>%
  select(-x) %>%
  left_join(thetas %>% filter(sample == spl)) %>%
  mutate(long = rbinom(nrow(.), 1, proba),
         value = ifelse(s == 1 & long == 0, delta, value)) %>%
    mutate(x = rexp(nrow(.), 1 / value)) %>%
    filter(x <= 40)
}

For each posterior, and for each observed sojourn in a given state, we draw a posterior prediction. We then only keep sojourns smaller than 40 minutes.

# We bin the observed holding times for plotting
pp_x_tbl <- post_pred_x %>%
  bind_rows() %>%
  # Binning sequence
  mutate(bin = sapply(x, function(val)
    seq(2.5, 37.5, by = 5)[which.min(abs(seq(2.5, 37.5, by = 5) - val))])) %>%
  group_by(s, sample, bin) %>%
  summarise(count = n()) %>%
  group_by(s, sample) %>%
  # Fill in the zeros
  complete(bin = seq(2.5, 37.5, by = 5), fill = list(count = 0)) %>%
  ungroup()

`summarise()` has grouped output by 's', 'sample'. You can override using the
`.groups` argument.

p1 <- obs_ht %>% filter(s == 1 & c == 0) %>%
hist_plot(
    colbin = "#c9c7bd",
    posterior = 1,
    post_data = pp_x_tbl %>%
      filter(s == 1),
    alpha_point = 0.6,
    alpha_line = 0.4,
    alpha_hist = 0.4,
    col_post = "#c9c7bd"
  )

Warning: Using `size` aesthetic for lines was deprecated in ggplot2 3.4.0.
ℹ Please use `linewidth` instead.

p2 <- obs_ht %>% filter(s == 2 & c == 0) %>%
hist_plot(
    colbin = "#E9C4C1",
    posterior = 1,
    post_data = pp_x_tbl %>%
      filter(s == 2),
    alpha_point = 0.6,
    alpha_line = 0.4,
    alpha_hist = 0.4,
    col_post = "#d1a7a3"
  )

p3 <- obs_ht %>% filter(s == 3 & c == 0) %>%
hist_plot(
    colbin = "#ce8da6",
    posterior = 1,
    post_data = pp_x_tbl %>%
      filter(s == 3),
    alpha_point = 0.6,
    alpha_line = 0.4,
    alpha_hist = 0.4,
    col_post = "#a37184"
  )

p4 <- obs_ht %>% filter(s == 4 & c == 0) %>%
hist_plot(
    colbin = "#996282",
    posterior = 1,
    post_data = pp_x_tbl %>%
      filter(s == 4),
    alpha_point = 0.6,
    alpha_line = 0.4,
    alpha_hist = 0.4,
    col_post = "#7a4e68"
  )

2.3.2 Transitions

obs_tr <- tibble(
dyad = d_1[[1]]$C_dyad,
s_from = d_1[[1]]$C_s_from
)

# gamma posterior samples
gammas <- post_1 %>%
  filter(iter == 1) %>%
  select(starts_with("Gamma[")) %>% 
  slice(1:30) %>%
  mutate(sample = 1:nrow(.)) %>%
  gather(param, value, 1:(ncol(.) - 1)) %>%
  extract(param, into = c("dyad", "s_from", "s_to"), 
          regex = "Gamma\\[(\\d+),(\\d+),(\\d+)\\]", convert = TRUE) %>%
  pivot_wider(names_from = s_to, 
              values_from = value, 
              names_prefix = "s_to_")


post_pred_tr <- list()
for (spl in 1:30) {
  post_pred_tr[[spl]] <- obs_tr %>%
  left_join(gammas %>% filter(sample == spl)) %>%
  mutate(s_to = mapply(function(s1, s2, s3, s4) sample(1:4, size = 1, prob = c(s1, s2, s3, s4)),
                      s_to_1, s_to_2, s_to_3, s_to_4))
  post_pred_tr[[spl]] <- post_pred_tr[[spl]] %>%
  group_by(s_from, s_to) %>%
  summarise(count = n()) %>%
    mutate(sample = spl)
}

Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.
Joining with `by = join_by(dyad, s_from)`
`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.

post_pred_tr_tbl <- post_pred_tr %>%
  bind_rows() %>%
  mutate(s_to = as.factor(s_to))

obs_tr_smry <- obs_tr %>%
  mutate(s_to = as.factor(d_1[[1]]$C_s_to)) %>%
  group_by(s_from, s_to) %>%
  summarise(count = n()) %>%
  ungroup()

`summarise()` has grouped output by 's_from'. You can override using the
`.groups` argument.

p5 <- binom_pmf(
  data = obs_tr_smry %>% filter(s_from == 1),
  col_data = c("#E9C4C1", "#ce8da6", "#996282"),
  col_post = c("#E9C4C1", "#ce8da6", "#996282"),
  posterior = 1,
  post_data = post_pred_tr_tbl %>% filter(s_from == 1)
)

p6 <- binom_pmf(
  data = obs_tr_smry %>% filter(s_from == 2),
  col_data = c("#c9c7bd", "#ce8da6", "#996282"),
  col_post = c("#ce8da6","#996282", "#c9c7bd"),
  posterior = 1,
  post_data = post_pred_tr_tbl %>% filter(s_from == 2)
)

p7 <- binom_pmf(
  data = obs_tr_smry %>% filter(s_from == 3),
  col_data = c("#c9c7bd", "#E9C4C1", "#996282"),
  col_post = c("#E9C4C1", "#996282", "#c9c7bd"),
  posterior = 1,
  post_data = post_pred_tr_tbl %>% filter(s_from == 3)
)

p8 <- binom_pmf(
  data = obs_tr_smry %>% filter(s_from == 4),
  col_data = c("#c9c7bd", "#E9C4C1", "#ce8da6"),
  col_post = c("#E9C4C1", "#ce8da6", "#c9c7bd"),
  posterior = 1,
  post_data = post_pred_tr_tbl %>% filter(s_from == 4)
)

wrap_plots(p1, p5, p2, p6, 
           p3, p7, p4, p8, ncol = 2)

2.4 Posterior contrasts

Target values from the generative simulation:

avgs <- list()
for (i in 1:1e3) {
  param_test <- stochastic_param(N_ind_vec = c(10, 10, 10))
   
  target <- list()
  for (gr in 1:3) {
    target[[gr]] <- tibble(
      z = paste0(param_test[[gr]]$dyads$z_a, "_", param_test[[gr]]$dyads$z_b),
      theta_1 = param_test[[gr]]$theta[, 1] / (12 * 60),
      theta_2 = param_test[[gr]]$theta[, 2],
      theta_3 = param_test[[gr]]$theta[, 3]
    )
  }
  
  avgs[[i]] <- target %>%
    bind_rows() %>%
    mutate(dyad = 1:nrow(.)) %>%
    group_by(z) %>%
    summarise(
      avg_1 = mean(theta_1),
      avg_2 = mean(theta_2),
      avg_3 = mean(theta_3),
      iter = i
    )
}

avgs <- avgs %>%
  bind_rows() %>%
  group_by(z) %>%
  summarise(
      avg_1 = mean(avg_1),
      avg_2 = mean(avg_2),
      avg_3 = mean(avg_3)
    )
avgs %>% saveRDS("./sim_data/02/param_avg.rds")

avgs <- readRDS("./sim_data/02/param_avg.rds")
p1 <- post_1 %>%
  select(`avg_s1[1]`, iter) %>%
  mutate(value = `avg_s1[1]` / (60 * 12)) %>%
  post_plot_2(
    min = 0,
    max = 20,
    lower_bound = 1,
    target_data = avgs$avg_1[1],
  )  +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank())

p2 <- post_1 %>%
  select(`avg_s1[2]`, iter) %>%
  mutate(value = `avg_s1[2]` / (60 * 12)) %>%
  post_plot_2(
    min = 0,
    max = 20,
    lower_bound = 1,
    y_space = -0.07,
    scale_dens = 0.5,
    target_data = avgs$avg_1[2],
  )  +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank()) +
  scale_y_continuous(limits = c(-0.142, 1.15),
                       expand = c(0, 0))

Scale for y is already present.
Adding another scale for y, which will replace the existing scale.

p3 <- post_1 %>%
  select(`avg_s1[3]`, iter) %>%
  mutate(value = `avg_s1[3]` / (60 * 12)) %>%
  post_plot_2(
    min = 0,
    max = 20,
    lower_bound = 1,
    target_data = avgs$avg_1[4],
  )

p4 <- post_1 %>%
  select(`avg_s2[1]`, iter) %>%
  mutate(value = `avg_s2[1]`) %>%
  post_plot_2(
    min = 0,
    max = 20,
    lower_bound = 1,
    target_data = avgs$avg_2[1],
  )  +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank())

p5 <- post_1 %>%
  select(`avg_s2[2]`, iter) %>%
  mutate(value = `avg_s2[2]`) %>%
  post_plot_2(
    min = 0,
    max = 20,
    lower_bound = 1,
    scale_dens = 0.5,
    y_space = -0.07,
    target_data = avgs$avg_2[2],
  ) +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank()) +
  scale_y_continuous(limits = c(-0.142, 1.15),
                       expand = c(0, 0))

Scale for y is already present.
Adding another scale for y, which will replace the existing scale.

p6 <- post_1 %>%
  select(`avg_s2[3]`, iter) %>%
  mutate(value = `avg_s2[3]`) %>%
  post_plot_2(
    min = 0,
    max = 20,
    lower_bound = 1,
    target_data = avgs$avg_2[4],
  )

p7 <- post_1 %>%
  select(`avg_gr[1]`, iter) %>%
  mutate(value = `avg_gr[1]`) %>%
  post_plot_2(
    min = 0,
    max = 20,
    lower_bound = 1,
    target_data = avgs$avg_3[1],
  ) +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank())

p8 <- post_1 %>%
  select(`avg_gr[2]`, iter) %>%
  mutate(value = `avg_gr[2]`) %>%
  post_plot_2(
    min = 0,
    max = 20,
    lower_bound = 1,
    target_data = avgs$avg_3[2],
  ) +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank())

p9 <- post_1 %>%
  select(`avg_gr[3]`, iter) %>%
  mutate(value = `avg_gr[3]`) %>%
  post_plot_2(
    min = 0,
    max = 20,
    lower_bound = 1,
    y_space = -0.175,
    target_data = avgs$avg_3[3],
  ) +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank())

p10 <- post_1 %>%
  select(`avg_gr[4]`, iter) %>%
  mutate(value = `avg_gr[4]`) %>%
  post_plot_2(
    min = 0,
    max = 20,
    lower_bound = 1,
    target_data = avgs$avg_3[4],
  )


design <-
"ADG
 BEH
 BEI
 CFJ"

p1 + p2 + p3 + p4 + p5 + 
p6 + p7 + p8 + p9 + p10 +
  plot_layout(design = design)

avgs <- readRDS("./sim_data/02/param_avg.rds")
p1 <- post_1 %>%
  select(`ATE_s1[1]`, iter) %>%
  mutate(value = `ATE_s1[1]` / (60 * 12)) %>%
  post_plot_2(
    min = -30,
    max = 20,
    y_space = -0.07,
    scale_dens = 0.5,
    target_data = avgs$avg_1[2] - avgs$avg_1[1],
  ) +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank()) +
  scale_y_continuous(limits = c(-0.142, 1.15),
                       expand = c(0, 0))

Scale for y is already present.
Adding another scale for y, which will replace the existing scale.

p2 <- post_1 %>%
  select(`ATE_s1[2]`, iter) %>%
  mutate(value = `ATE_s1[2]` / (60 * 12)) %>%
  post_plot_2(
    min = -30,
    max = 20,
    target_data = avgs$avg_1[4] - avgs$avg_1[1],
  )

p3 <- post_1 %>%
  select(`ATE_s2[1]`, iter) %>%
  rename(value = `ATE_s2[1]`) %>%
  post_plot_2(
    min = -10,
    max = 15,
    y_space = -0.07,
    scale_dens = 0.5,
    target_data = avgs$avg_2[2] - avgs$avg_2[1],
  ) +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank()) +
  scale_y_continuous(limits = c(-0.142, 1.15),
                       expand = c(0, 0))

Scale for y is already present.
Adding another scale for y, which will replace the existing scale.

p4 <- post_1 %>%
  select(`ATE_s2[2]`, iter) %>%
  rename(value = `ATE_s2[2]`) %>%
  post_plot_2(
    min = -10,
    max = 15,
    target_data = avgs$avg_2[4] - avgs$avg_2[1],
  )

p5 <- post_1 %>%
  select(`ATE_gr[1]`, iter) %>%
  rename(value = `ATE_gr[1]`) %>%
  post_plot_2(
    min = -10,
    max = 15,
    target_data = avgs$avg_3[2] - avgs$avg_3[1],
  ) +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank())

p6 <- post_1 %>%
  select(`ATE_gr[2]`, iter) %>%
  rename(value = `ATE_gr[2]`) %>%
  post_plot_2(
    min = -10,
    max = 15,
     target_data = avgs$avg_3[3] - avgs$avg_3[1],
  )  +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank())

p7 <- post_1 %>%
  select(`ATE_gr[3]`, iter) %>%
  rename(value = `ATE_gr[3]`) %>%
  post_plot_2(
    min = -10,
    max = 15,
     target_data = avgs$avg_3[4] - avgs$avg_3[1],
  )

design <-
"ACE
 ACF
 BDG"

p1 + p2 + p3 + 
p4 + p5 + p6 + p7 +
  plot_layout(design = design)